# vim: expandtab:ts=4:sw=4
from __future__ import absolute_import
import numpy as np
from . import kalman_filter
from . import linear_assignment
from . import iou_matching
from .track import Track
import configuration as cfg

class Tracker:
    """
    这是一个多目标跟踪器

    Parameters
    ----------
    metric : nn_matching.NearestNeighborDistanceMetric
        一个 测量到跟踪关联的 距离度量
    Attributes
    ----------
    max_age : int
        在一个track被删除之前，未命中的最大值
    n_init : int
        确认轨迹前的连续检测次数。如果在前"n_init"帧内发生未命中，则跟踪状态置为"Deleted"
    kf : kalman_filter.KalmanFilter
        卡尔曼滤波器
    tracks : List[Track]
        当前step的active tracks
    """
    def __init__(self, metric):
        self.metric = metric    #用于衡量跟踪关联的距离度量,cos或者欧几里得距离
        self.kf = kalman_filter.KalmanFilter()  #卡尔曼滤波器
        self.tracks = []    #当前时刻活动的track列表，保存着所有的tracks
        self._next_id = 1

    def predict(self):
        """
        向前传播跟踪状态分布
        这个函数应该每个step都被调用一次，在update()之前
        """
        #对每一条track使用卡尔曼滤波的predict
        for track in self.tracks:
            track.predict(self.kf)

    def update(self, detections,class_ids):
        """
        执行测量更新和跟踪管理，将检测结果detections与跟踪结果tracks进行匹配
        ----------
        detections : List[deep_sort.detection.Detection]
            当前时刻的检测列表
        """
        #匹配级联+IOU匹配，完成新增detections和现有tracks的匹配
        # 返回匹配到的tracks，未匹配到的tracks，未匹配到的detections
        matches, unmatched_tracks, unmatched_detections = self._match(detections)

        # 对于所有匹配到的track，使用detections来更新track
        for track_idx, detection_idx in matches:
            self.tracks[track_idx].update(self.kf, detections[detection_idx],class_ids[detection_idx])
        '''
        下面是未匹配到的两种现象，可能是未匹配到tracks，
        也有可能是未匹配到的detections
        --------------------------------------------------
        对于未匹配到的tracks，可能是检测器捡漏了目标，
        或者目标在视野里丢失，此时查看其是否是confirmed，
        如果是unconfirmed则将其删除，反之为其设置寿命age，
        在寿命的时间内(>max_age)，其状态仍然是unconfirmed，则删除，
        在寿命之内(<max_age)，对其进行三次考察，决定是否修改其状态
        --------------------------------------------------
        对于未匹配到的detections，可能是新出现的目标，可能是
        长时间遮挡，此时为其建立新tracks，
        并将状态设置为unconfirmed，对其进行三次考察，如果是目标，则
        修改为confirmed
        '''
        for track_idx in unmatched_tracks:#未匹配到的tracks
            self.tracks[track_idx].mark_missed()

        for detection_idx in unmatched_detections:#未匹配到的detections
            self._initiate_track(detections[detection_idx],class_ids[detection_idx].item())

        self.tracks = [t for t in self.tracks if not t.is_deleted()]

        #更新距离度量metric
        active_targets = [t.track_id for t in self.tracks if t.is_confirmed()]
        features, targets = [], []
        for track in self.tracks:
            if not track.is_confirmed():
                continue
            features += track.features
            targets += [track.track_id for _ in track.features]
            track.features = []
        self.metric.partial_fit(
            np.asarray(features), np.asarray(targets), active_targets)

    def _match(self, detections):
        '''
        将track与detections进行级联匹配
        detections:是一个list，保存着经过NMS的bounding box，包含置信度，tlwh坐标，和box的特征
        '''

        def gated_metric(tracks, dets, track_indices, detection_indices):
            '''
            根据外观信息和马氏距离，计算卡尔曼滤波预测到的tracks和当前时刻检测到的
            detections的代价矩阵
            '''
            features = np.array([dets[i].feature for i in detection_indices])
            targets = np.array([tracks[i].track_id for i in track_indices])
            #基于外观信息，计算tracks和detections的余弦距离代价矩阵
            cost_matrix = self.metric.distance(features, targets)
            #基于马氏距离，过滤掉代价矩阵中的一些不合理的项，将其设置为一个较大的值
            cost_matrix = linear_assignment.gate_cost_matrix(
                self.kf, cost_matrix, tracks, dets, track_indices,
                detection_indices)
            return cost_matrix

        #将tracks集合拆分为已确认和未确认，只有confirmed的tracks会进行匹配
        confirmed_tracks = [i for i, t in enumerate(self.tracks) if t.is_confirmed()]
        unconfirmed_tracks = [i for i, t in enumerate(self.tracks) if not t.is_confirmed()]

        #使用外观特征（余弦距离或者马氏距离）关联所有confirmed的tracks,返回匹配结果，
        #matched,unmatched_tracks,unmatched_detections。
        matches_a, unmatched_tracks_a, unmatched_detections = \
            linear_assignment.matching_cascade(
                distance_metric=gated_metric,
                tracks=self.tracks,
                detections=detections,
                track_indices=confirmed_tracks)

        '''
        使用IOU匹配将剩余的unmatched_tracks的time_since_update=1
        和unconfirmed_tracks的tracks进行匹配。
        发生在级联匹配之后，对象为上一轮级联匹配失败中age=1的tracks，
        有助于解决因上一帧的部分遮挡而引起的突然出现的外观变化，从而
        减小被遗漏的概率
        '''
        #从上一轮级联匹配失败的跟踪器中挑选出连续1帧没有匹配上（相当于age=1）
        #的跟踪器，并和unconfirmed_tracks相加。
        iou_track_candidates = unconfirmed_tracks + [k for k in unmatched_tracks_a if
                                                 self.tracks[k].time_since_update == 1]
        unmatched_tracks_a = [k for k in unmatched_tracks_a if
                              self.tracks[k].time_since_update != 1]
        #将tracks与剩下没匹配上的detections进行IOU匹配
        matches_b, unmatched_tracks_b, unmatched_detections = \
            linear_assignment.min_cost_matching(
                iou_matching.iou_cost, self.tracks,
                detections, iou_track_candidates, unmatched_detections)

        matches = matches_a + matches_b
        unmatched_tracks = list(set(unmatched_tracks_a + unmatched_tracks_b))
        return matches, unmatched_tracks, unmatched_detections

    def _initiate_track(self, detection,class_id):
        means, covariances = self.kf.initiate(detection.to_xyah())
        self.tracks.append(Track(
            means, covariances, self._next_id, cfg.n_init, cfg.max_age,class_id,
            detection.feature))
        self._next_id += 1
